'''
灰色的羊吃草，黑色的狼吃羊。他们能维持生态的平衡吗？

作者：侯展意
协议：木兰2.0
'''
import numpy as np
import takagiabm as tak  # 这里分别用了两种导入风格。


def createSheep():  # 调用这个函数生成一只“羊”
    a = tak.GridAgent()
    triangle2 = tak.scale(tak.takTriangle, 3)  # 设置形状为takagi库的三角形，缩放比例为3.
    a.setTrait('energy', 10)
    a.setShape(triangle2)  # 设置形状为三角形。
    a.setSpeed(np.array([1, 0]))  # 设置初始速度
    a.setColor('#aaaaaa')  # 设置颜色，浅灰色接近白色
    a.setTrait('type', 'sheep')
    return a


def createWolf():  # 生成一只狼。

    a = tak.GridAgent()
    triangle2 = tak.scale(tak.takTriangle, 3)  # 设置形状为takagi库的三角形，缩放比例为3.
    a.setTrait('energy', 40)
    a.setTrait('type', 'wolf')
    a.setShape(triangle2)  # 设置形状为三角形。
    a.setSpeed(np.array([1, 0]))  # 设置初始速度
    a.setColor('#666666')  # 设置颜色,深灰色接近黑色
    return a


def agentSetup(model):
    for i in range(20):
        a = createWolf()
        model.addAgent(a, randomPlace=True)
    for i in range(40):
        a = createSheep()
        model.addAgent(a, randomPlace=True)
    for i in range(model.width):
        for j in range(model.height):
            if (np.random.random() < 0.5):
                model.grid.cells[j][i].setTrait('color', '#00ff00')


def wolfAction(agent: tak.GridAgent):
    global wolfBirth
    agentList = agent.getAgents()
    if agentList != []:
        for a in agentList:
            if (a.traits['type'] == 'sheep'):
                agent.kill(a)  # 把羊吃了
                agent.traits['energy'] += a.traits['energy']  # 获取羊的能量

    if (np.random.random() < wolfBirth.value):  # 取wolfBirth的值。
        if (agent.getTrait('energy') > 10):
            a = createWolf()
            a.setTrait('energy', 10)

            agent.setTrait('energy', agent.getTrait('energy') - 10)

            a.setPos(agent.pos.astype(np.float))
            agent.model.addAgent(a)

    energy = agent.getTrait('energy')
    agent.setTrait('energy', energy - 1)
    if (agent.getTrait('energy') <= 0):
        agent.die()
        return
    agent.turn(np.random.random() * np.pi - np.pi / 2)
    agent.move()


def sheepAction(agent):
    global sheepBirth
    color = agent.getCellColor()
    if (color == '#00ff00'):
        agent.setCellColor('#ffffff')
        agent.properties['energy'] += 5

    if (np.random.random() < sheepBirth.value):
        if (agent.getTrait('energy') > 10):
            a = createSheep()
            a.setTrait('energy', 5)
            agent.setTrait('energy', agent.getTrait('energy') - 5)
            a.setPos(agent.pos.astype(np.float))
            agent.model.addAgent(a)

    agent.traits['energy'] -= 1
    if (agent.getTrait('energy') <= 0):
        agent.die()
        return
    agent.turn(np.random.random() * np.pi - np.pi / 2)
    agent.move()


def agentLoop(agent):  # 代理人(也就是羊)执行的函数
    global stat, sheepGrowth
    if (agent.traits['type'] == 'sheep'):
        sheepAction(agent)
    else:
        wolfAction(agent)


def cellStep(cell):  # 草的生长。
    global grassGrowth
    if (np.random.random() < grassGrowth.value):
        if (cell.getColor() == '#ffffff'):
            cell.setColor('#00ff00')


def cellOnClicked(cell):
    if (cell.getColor() == '#ffffff'):
        cell.setColor('#00ff00')
    else:
        cell.setColor('#ffffff')

sheepBirth = tak.Var('羊出生率', 0.03, range=(0, 0.01, 0.05))
grassGrowth = tak.Var('草生长率', 0.05, range=(0, 0.01, 0.15))
wolfBirth = tak.Var('狼出生率', 0.03, range=(0, 0.01, 0.05))

varList = [grassGrowth, sheepBirth, wolfBirth]

dataCounterList = [tak.DataCounter(name='颜色', propertyName='color', targets=['cell']),
                   tak.DataCounter(propertyName='color', targets=['agent'])]
tak.GridAgent = tak.prepareAgent(tak.GridAgent, agentStepFunc=agentLoop)
tak.Cell = tak.prepareCell(tak.Cell, cellStepFunc=cellStep)
tak.GridModel = tak.prepareModel(tak.GridModel, varList=varList, dataCounterList=dataCounterList
                                 , width=50, height=50,modelInitFunc=agentSetup)

if __name__ == "__main__":
    tak.simStart(__file__, tak.GridModel, maxSteps=10000)  # 这个方法同时启动了图形界面。
